// Copyright DustedPixels.com 2008. All rights reserved.

package com.dustedpixels.asm;

import org.objectweb.asm.ClassAdapter;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.AdviceAdapter;
import org.objectweb.asm.commons.RemappingMethodAdapter;
import org.objectweb.asm.commons.SimpleRemapper;

import com.dustedpixels.common.base.Repeatable;

/**
 * @author micapolos@gmail.com (Michal Pociecha-Los)
 */
public final class RepeatableClassVisitor extends ClassAdapter {

  private final static String OBJECT_TYPE_NAME = Type.getInternalName(Object.class);
  private final static String RUNNABLE_TYPE_NAME = Type.getInternalName(Runnable.class);
  private final static String REPEATABLE_TYPE_NAME = Type.getInternalName(Repeatable.class);

  private final static String RUN_METHOD_NAME = "run";
  private final static String RUN_METHOD_DESC = 
      Type.getMethodDescriptor(Type.VOID_TYPE, new Type[0]);
  
  private final static String REPEAT_METHOD_NAME = "repeat";
  private final static String REPEAT_METHOD_DESC =
      Type.getMethodDescriptor(Type.VOID_TYPE, new Type[] {Type.INT_TYPE});
  
  private String oldClassName;
  private final String newClassName;

  public RepeatableClassVisitor(ClassVisitor out, String className) {
    super(out);
    this.newClassName = className;
  }

  @Override
  public void visit(int version, int access, String name, String signature, String superName,
      String[] interfaces) {
    
    if (!superName.equals(OBJECT_TYPE_NAME)) {
      throw new RuntimeException("Must inherit from " + OBJECT_TYPE_NAME);
    }
    
    int runnableIndex = -1;
    for (int i = 0; i < interfaces.length; i++) {
      if (interfaces[i].equals(RUNNABLE_TYPE_NAME)) {
        runnableIndex = i;
        break;
      }
    }
    
    if (runnableIndex == -1) {
      throw new RuntimeException("Must implement " + RUNNABLE_TYPE_NAME);
    }
    
    String[] newInterfaces = new String[interfaces.length];
    System.arraycopy(interfaces, 0, newInterfaces, 0, interfaces.length);
    newInterfaces[runnableIndex] = REPEATABLE_TYPE_NAME;
    
    this.oldClassName = name;
    
    super.visit(version, access, newClassName, signature, superName, newInterfaces);
  }

  @Override
  public MethodVisitor visitMethod(int access, String name, String desc, String signature,
      String[] exceptions) {
    
    MethodVisitor mv;
    
    if (name.equals(RUN_METHOD_NAME) && desc.equals(RUN_METHOD_DESC)) {
      mv = super.visitMethod(access, REPEAT_METHOD_NAME, REPEAT_METHOD_DESC, signature, exceptions);
      mv = new RepeatableMethodVisitor(mv, access, REPEAT_METHOD_NAME, REPEAT_METHOD_DESC);
    } else {
      mv = super.visitMethod(access, name, desc, signature, exceptions);
    }

    mv = new RemappingMethodAdapter(access, REPEAT_METHOD_DESC, mv, new SimpleRemapper(oldClassName, newClassName));
    
    return mv;
  }

  private static final class RepeatableMethodVisitor extends AdviceAdapter {
    private final Label startLabel;
    private final Label endLabel;
    
    protected RepeatableMethodVisitor(MethodVisitor mv, int access, String name, String desc) {
      super(mv, access, name, desc);
      this.startLabel = new Label();
      this.endLabel = new Label();
    }

    @Override
    public void visitVarInsn(int opcode, int var) {
      if (var != 0) var++;
      super.visitVarInsn(opcode, var);
    }

    @Override
    public void visitIincInsn(int var, int increment) {
      if (var != 0) var++;
      super.visitIincInsn(var, increment);
    }

    @Override
    protected void onMethodEnter() {
      super.visitJumpInsn(Opcodes.GOTO, endLabel);
      super.visitLabel(startLabel);
    }

    @Override
    protected void onMethodExit(int opcode) {
      super.visitLabel(endLabel);
      super.visitVarInsn(Opcodes.ILOAD, 1);
      super.visitIincInsn(1, -1);
      super.visitJumpInsn(Opcodes.IFGT, startLabel);
    }
  }
}
